// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/functions/sparksql/unscaled_value_function.h>

#include <pollux/expression/decoded_args.h>

namespace kumo::pollux::functions::sparksql {
    namespace {
        // Return the unscaled bigint value of a decimal, assuming it
        // fits in a bigint. Only short decimal input is accepted.
        class UnscaledValueFunction final : public exec::VectorFunction {
            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args,
                const TypePtr &outputType,
                exec::EvalCtx &context,
                VectorPtr &result) const final {
                POLLUX_USER_CHECK(
                    args[0]->type()->isShortDecimal(),
                    "Expect short decimal type, but got: {}",
                    args[0]->type());
                exec::DecodedArgs decodedArgs(rows, args, context);
                auto decimalVector = decodedArgs.at(0);
                context.ensure_writable(rows, BIGINT(), result);
                result->clear_nulls(rows);
                auto flatResult =
                        result->as_unchecked<FlatVector<int64_t> >()->mutableRawValues();
                rows.applyToSelected([&](auto row) {
                    flatResult[row] = decimalVector->value_at<int64_t>(row);
                });
            }
        };
    } // namespace

    std::vector<std::shared_ptr<exec::FunctionSignature> >
    unscaledValueSignatures() {
        return {
            exec::FunctionSignatureBuilder()
            // precision <= 18.
            .integerVariable("precision", "min(precision, 18)")
            .integerVariable("scale")
            .returnType("bigint")
            .argumentType("DECIMAL(precision, scale)")
            .build()
        };
    }

    std::unique_ptr<exec::VectorFunction> makeUnscaledValue() {
        return std::make_unique<UnscaledValueFunction>();
    }
} // namespace kumo::pollux::functions::sparksql
